import numpy as np
import random
import copy
from Init_population import GA

class Individual:

    def __init__(self, obj_index, config):
        
        self.config = config
        self.obj_index = obj_index
        self.x = np.random.randint(low=-1, high=2, size=self.config['DsDim'])
        self.fitness = self.config['evaluationer'].eval(self.x)
        self.distance = 0.0

    def update(self, solution1, solution2, solution3):

        ## cross1
        self.cross_x = np.zeros_like(self.x)
        for index in range(len(solution1)):
            r = random.random()
            if r < self.config['cross_rate']:
                self.cross_x[index] = solution1[index]
            else:
                self.cross_x[index] = solution2[index]
        
        ## cross2, solution3 comes from archive
        for index in range(len(solution3)):
            r = random.random()
            if r < self.config['cross_rate']:
                self.cross_x[index] = solution3[index]

        ## mutation
        for index in range(len(self.cross_x)):
            r = random.random()
            if r < self.config['mutation_rate']:
                temp = np.random.randint(low=-1, high=2)
                while temp == self.cross_x[index]:
                    temp = np.random.randint(low=-1, high=2)
                self.cross_x[index] = temp
        self.x = self.cross_x
        self.fitness = self.config['evaluationer'].eval(self.x)
    
    
    def deepcopy(self, copyed_individual):

        self.obj_index = copyed_individual.obj_index
        self.x = copy.deepcopy(copyed_individual.x)
        self.fitness = copy.deepcopy(copyed_individual.fitness)
    
    def els(self, ):
        
        pos = random.randint(0, len(self.x) - 1)
        temp = np.random.randint(low=-1, high=2)
        while temp == self.x[pos]:
            temp = np.random.randint(low=-1, high=2)
        self.x[pos] = temp
        self.fitness = self.config['evaluationer'].eval(self.x)

    def reset_distance(self, ):

        self.distance = 0.0


class CMEA:

    def __init__(self, config):

        self.config = config
        self.OsDim = self.config['OsDim']
        self.gaer = GA(self.config)

        self.populations = []

        ## init populations
        for obj_index in range(self.OsDim):
            
            population = [Individual(obj_index=obj_index, config=self.config) for i in range(self.config['NP'])]
            self.populations.append(population)
        
        ##init pBest population
        self.pBest = []
        for obj_index in range(self.OsDim):
            pbest = []
            for i in range(self.config['NP']):
                temp = Individual(obj_index=obj_index, config=self.config)
                temp.deepcopy(self.populations[obj_index][i])
                pbest.append(temp)
            self.pBest.append(pbest)


        self.archive = {
            'upper_size': self.config['arc_size'],
            'size': 0,
            'set': [],
        }

    def init_solution(self, ):

        self.gaer.update()
        self.gaer.show()
    
    def init(self, ):

        count = 0
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                self.populations[obj_index][i].x = self.gaer.population[count].x
                count += 1
                self.populations[obj_index][i].fitness = self.config['evaluationer'].eval(self.populations[obj_index][i].x)

    def update_pBest(self, ):

        for obj_index in range(self.OsDim):
            
            for i in range(self.config['NP']):

                if self.populations[obj_index][i].fitness[obj_index] < self.pBest[obj_index][i].fitness[obj_index]:

                    self.pBest[obj_index][i].deepcopy(self.populations[obj_index][i])

    def update_archive(self, ):
        
        def dominate(fitness_1, fitness_2):
            """
            return 0/-1/1, where 0 represent co-dominate, -1 represent fit_2 domainte fit_1, 1 represent fit 1 domainate fit_2.
            """
            better = (fitness_1 < fitness_2).sum()
            same = (fitness_1 == fitness_2).sum()
            worse = (fitness_1 > fitness_2).sum()

            if same == fitness_1.shape[0]:
                return 0
            elif better == 0:
                return -1
            elif worse == 0:
                return 1
            else:
                return 0

        def nondominatedSet(S_set):
            """_summary_
                Get non-dominated Set R and return.
            Args:
                S_set (_type_): _description_
            Retrun:
                R_set (List):
            """
            size = len(S_set)
            R_set = []

            for i in range(size):
                non = True
                individual_a = S_set[i]
                for j in range(size):
                    if i != j:
                        individual_b = S_set[j]
                        k = dominate(individual_a.fitness, individual_b.fitness)
                        if k < 0:
                            non = False
                            break
                if non:
                    R_set.append(individual_a)
            return R_set            

        def crowding_distance_assignment(R_set):
            """_summary_
                Give each individuals a crowding distance, and return the "sorted(by distance)" set.
            Args:
                R_set (List): 
            """
            ## zero the distance for individuals in R_set
            size = len(R_set)
            for individual in R_set:
                individual.reset_distance()
            for sort_obj_index in range(self.OsDim):
                # sort by the i-th objective
                R_set.sort(key=lambda individual: individual.fitness[sort_obj_index])
                R_set[0].distance = 1e6
                R_set[size - 1].distance = 1e6
                fitmin = R_set[0].fitness[sort_obj_index]
                fitmax = R_set[size - 1].fitness[sort_obj_index]

                for index in range(1, size - 1):
                    R_set[index].distance += (R_set[index + 1].fitness[sort_obj_index] - R_set[index - 1].fitness[sort_obj_index]) / (fitmax - fitmin)
            
            # sort by distance
            R_set.sort(key=lambda individual: individual.distance, reverse=True)

            return R_set

        S = [] 
        
        ## Add pBest of each individual in each swarm into set s
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                S.append(self.pBest[obj_index][i])


        ## Add all the solution in the old archive into the set S
        for i in range(self.archive['size']):
            S.append(self.archive['set'][i])
        

        ## ELS strategy
        for i in range(self.archive['size']):
            temp = Individual(obj_index=self.archive['set'][i].obj_index, config=self.config)
            temp.deepcopy(self.archive['set'][i])
            temp.els()
            S.append(temp)
        

        R = nondominatedSet(S)


        if len(R) < self.archive['upper_size']:
            self.archive['set'] = copy.deepcopy(R)
            
            if len(self.archive['set']) != len(R):
                raise ValueError
            
            self.archive['size'] = len(self.archive['set'])
        else:
            sorted_R = crowding_distance_assignment(R)
            self.archive['set'].clear()
            for i in range(self.archive['upper_size']):
                self.archive['set'].append(sorted_R[i])
            if len(self.archive['set']) != self.archive['upper_size']:
                raise ValueError
            self.archive['size'] = len(self.archive['set'])    

            del sorted_R
        
        del S 
        del R   

    def update_populations(self, ):

        self.solutions = np.zeros((self.OsDim, self.config['NP'], 3, self.config['DsDim']))

        for obj_index in range(self.OsDim):
        
            count = 0
            while count < 2 * self.config['NP']:

                lister = np.random.permutation(self.config['NP'])
                for index in range(0, self.config['NP'], 2):
                    left = lister[index]
                    right = lister[index + 1]
                    if self.populations[obj_index][left].fitness[obj_index] < self.populations[obj_index][right].fitness[obj_index]:
                        self.solutions[obj_index, count // 2, count % 2] = self.populations[obj_index][left].x
                    else:
                        self.solutions[obj_index, count // 2, count % 2] = self.populations[obj_index][right].x
                    count += 1
            
            if self.archive['size'] == 0:
                for i in range(self.config['NP']):
                    pos = random.randint(0, self.config['NP'] - 1)
                    self.solutions[obj_index, i, 2] = self.pBest[(obj_index + 1) % 2][pos].x
            else:
                for i in range(self.config['NP']):
                    pos = random.randint(0, self.archive['size'] - 1)
                    self.solutions[obj_index, i, 2] = self.archive['set'][pos].x

        for obj_index in range(self.OsDim):
            for index in range(self.config['NP']):
                self.populations[obj_index][index].update(self.solutions[obj_index, index, 0], self.solutions[obj_index, index, 1], self.solutions[obj_index, index, 2])

        self.update_pBest()
        
    def show(self, ):

        for individual in self.population:
            print(individual.fitness)